[Fix] Reject non-divisible GQA head counts - #1032

Summary

Several attention paths derived the GQA group size using HQ // H without first requiring divisibility. A nonzero remainder can leave query heads uncovered or map a query head beyond the available KV heads.

Add a shared group-size validator and use it across attention, decoding, NSA, Parallax, DSA, and forgetting-attention reference paths. NSA retains its additional power-of-two and minimum group-size requirements.

Valid divisible layouts are unchanged. Invalid layouts now raise ValueError before repetition or kernel launch.

Fixes #1029.

Failure modes

Integer division gives HQ = H * G + r, where G = HQ // H and r = HQ % H. Equal GQA groups cover exactly H * G query heads, so any r > 0 necessarily leaves a remainder outside that mapping.

NSA demonstrates the uncovered-head case:

HQ = 33, H = 2, G = 16

KV head 0  -> query heads  0..15
KV head 1  -> query heads 16..31
uncovered  -> query head      32

The old NSA guard accepted G=16 because it is a power of two and at least 16.

Query-head-indexed kernels can instead produce an invalid KV index:

HQ = 3, H = 2, G = 1

query head 0 -> KV head 0
query head 1 -> KV head 1
query head 2 -> KV head 2  (out of range)

The shared helper rejects H == 0 and HQ % H != 0 before either mapping can occur.

Scope

The same invariant is applied to optimized and reference paths so test/reference behavior matches public entrypoints. No kernel math or valid dispatch behavior changes.

Suggested review order

1. fla/ops/utils/head.py
2. Public attention, NSA, and Parallax entrypoints
3. Internal compression, decode, and reference paths
4. Invalid-layout regression coverage

Test plan

• invalid ratios are covered in attention, decoding, NSA, Parallax, and forgetting-attention tests
• python3 -m compileall -q fla/ops/attn fla/ops/nsa fla/ops/parallax fla/ops/dsa/naive.py fla/ops/forgetting_attn/naive.py fla/ops/utils/head.py
• ruff check fla/ops/attn fla/ops/nsa fla/ops/parallax fla/ops/dsa/naive.py fla/ops/forgetting_attn/naive.py fla/ops/utils/head.py
• affected-file pre-commit hooks

Compatibility

This changes only previously invalid GQA layouts. Valid callers are unaffected.